package cn.zwx.common.limiter.service;

import cn.zwx.base.exception.BizWrapException;
import cn.zwx.common.limiter.annotation.LimitRule;
import cn.zwx.common.limiter.enums.LimiterTypeEnum;
import com.google.common.util.concurrent.RateLimiter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import java.time.Duration;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author zhangwenxue
 */
@Service
public class MemoryLimiterServiceImpl implements LimiterService {

    public static final Logger logger = LoggerFactory.getLogger(MemoryLimiterServiceImpl.class);

    private final Map<String, RateLimiter> rateLimiters = new ConcurrentHashMap<>(16);


    @Override
    public String getType() {
        return LimiterTypeEnum.memory.getMsg();
    }

    @Override
    public boolean take(LimitRule limitRule) throws BizWrapException {
        String cacheKey =  limitRule.namespace() + limitRule.key();
        RateLimiter rateLimiter = rateLimiters.get(cacheKey);
        if (Objects.isNull(rateLimiter)){
            rateLimiter = RateLimiter.create(limitRule.quota(), Duration.ofSeconds(limitRule.period()));
            rateLimiters.putIfAbsent(cacheKey,rateLimiter);
        }
        if (rateLimiter.tryAcquire()){
            logger.info("得到令牌，准备执行业务 start time:{}",System.currentTimeMillis());
            return true;
        }
        return false;
    }

}
